# 尝试实现异步的训练过程？需要吗？暂时还不需要。
from asyncio.tasks import wait
import os 
if os.name == 'nt':
    os.chdir("C:\\federated_malware\\Project")
    print ("cwd",os.getcwd())#获得当前目录
    print ("工作目录",os.path.abspath('.'))#获得当前工作目录
    print ("工作目录",os.path.abspath(os.curdir))#获得当前工作目录
import logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# 导入参数------------------------------------------------------
from fl_config import dataset_params,server_params,client_params

logger.info(dataset_params)
logger.info(server_params)
logger.info(client_params)


# 导入数据------------------------------------------------------
# 初始化各种类型的数据集。这里返回的是不同数据的索引，不能修改，否则会影响原来的数据。可以复用这些索引。因为创建dataset的时候一般是复制一份，而且一般也不会对原来的数据集进行修改。
from datasets.minist.data_load import dataManager as mdm


# 只需要能跑过就行，统计的返回值Counter符合要求即可.
# 确实，规定数据量不太合理。导致数据缺失，显然结果会不好。
# 所以现在完全采取互斥划分的五种方式。正好可以减少数据划分的情况。
mdm_avg = mdm.allocate_data_avg(client_number=dataset_params["client_number"])
mdm_ln = mdm.allocate_data_noniid_by_label_number(client_number=dataset_params["client_number"],class_total=10,class_counts=dataset_params["class_counts"])
mdm_ld = mdm.allocate_data_noniid_by_label_distribution(client_number=dataset_params["client_number"],class_total=10,beta=dataset_params["beta"])
mdm_fn = mdm.allocate_data_noniid_by_feature_noise(client_number=dataset_params["client_number"],bias_level=dataset_params["bias_level"])
mdm_quantity = mdm.allocate_data_noniid_by_quantity(client_number=dataset_params["client_number"],beta=dataset_params["beta"])
mdm_test = mdm.allocate_data_test()

# 导入模型------------------------------------------------------
from models.minist.dnn_model import Net
# from models.maldroid.dnn_model import Net 
model = Net()
# param_list = model.parameters()

# 导入算法------------------------------------------------------
from trainers.fedbase.FedServerBase import FedServerBase
from trainers.fedbase.FedClientBase import FedClientBase
from trainers.fedavg.FedAvgServer import FedAvgServer
from trainers.fedavg.FedAvgClient import FedAvgClient

# 导入rpc---------------------------------------------------------
from rpc import trainer_service,trainer_proxy
import numpy as np
import asyncio
# ------------------------------------------------fedbase----------------------------------------
# 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。
train_dataset_list = [mdm_avg]
test_dataset_list = [mdm_test]
dataset_name = ['mdm_avg']


base_name = './results/'
algorithm_name = '_fedbase'
suffix = '_.pt'

for name,train_data,test_data in zip(dataset_name,train_dataset_list,test_dataset_list):

    # 创建多个rpc客户端
    train_worker_list = [FedClientBase(train_dataset,model) for train_dataset in train_data]
    test_worker_list = [FedClientBase(test_dataset,model) for test_dataset in test_data]

    # 为每个客户端生成端口号
    port_list = np.linspace(5000,6000,num=20,dtype=int)
    worker_id_list = []
    for prot in port_list:
        worker_id_list.append("127.0.0.1:{:d}".format(prot))
    
    port_list2 = np.linspace(7000,8000,num=1,dtype=int)
    test_id_list = []
    for prot in port_list2:
        test_id_list.append("127.0.0.1:{:d}".format(prot))
    
    
    # 通过异步的方式启动这些客户端，不能说是添加一层代理。利用这些客户端提供服务
    async def run_client():
        task = []
        # loop = asyncio.get_event_loop()
        for worker,worker_id in zip(train_worker_list,worker_id_list):
            # task.append(asyncio.create_task(run_test(worker_id)))
            task.append(trainer_service.start_server_async(worker,worker_id))

        for worker,worker_id in zip(test_worker_list,test_id_list):
            # task.append(asyncio.create_task(run_test(worker_id)))
            task.append(trainer_service.start_server_async(worker,worker_id))
        
        await asyncio.wait(task)

  
    asyncio.run(run_client())
    print('wait for finish')





# ------------------------------------------------fedavg----------------------------------------
# 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。

# train_dataset_list = [mdm_avg,mdm_ln,mdm_ld,mdm_fn,mdm_quantity]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_ln','mdm_ld','mdm_fn','mdm_quantity']
# train_dataset_list = [fdm_noniid]
# test_dataset_list = [fdm_test]
# dataset_name = ['fdm_noniid']

# train_dataset_list = [maldm_avg,maldm_ln,maldm_ld,maldm_fn,maldm_quantity]
# test_dataset_list = [maldm_test,maldm_test,maldm_test,maldm_test,maldm_test]
# dataset_name = ['maldm_avg','maldm_ln','maldm_ld','maldm_fn','maldm_quantity']

# base_name = './results/'
# algorithm_name = '_fedavg'
# suffix = '_.pt'

# for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
#     train_worker_list = [FedAvgClient(train_dataset,model) for train_dataset in train_data]
#     test_worker_list = [FedAvgClient(test_dataset,model) for test_dataset in test_dataset]
#     server = FedAvgServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
    
#     logger.info("本次测试的数据集的名字:{}".format(name))
#     logger.info("本次测试的算法的名字:{}".format(algorithm_name))
#     server.start_federated()
#     server.save_state(base_name+name+algorithm_name+suffix)

